package net.roseboy.classfinal;

import a.a.a.B;
import a.a.a.C;
import javassist.ClassPool;
import javassist.CtClass;
import net.roseboy.classfinal.util.*;

import java.io.ByteArrayInputStream;
import java.io.File;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

import static net.roseboy.classfinal.util.StrUtils.toBytes;

/**
 * java class加密
 *
 * @author roseboy
 */
public class JarEncryptor {
    //加密配置文件：加载配置文件是注入解密代码的配置
    static Map<String, String> aopMap = new HashMap<>();

    static {
        //org.springframework.core.io.ClassPathResource#getInputStream注入解密功能
        aopMap.put("spring.class", "org.springframework.core.io.ClassPathResource#getInputStream");
        aopMap.put("spring.code", "char[] c=${passchar};"
                + "is=net.roseboy.classfinal.JarDecryptor.getInstance().decryptConfigFile(this.path,is,c);");
        aopMap.put("spring.line", "999");

        //com.jfinal.kit.Prop#getInputStream注入解密功能
        aopMap.put("jfinal.class", "com.jfinal.kit.Prop#<Prop>(java.lang.String,java.lang.String)");
        aopMap.put("jfinal.code", "char[] c=${passchar};inputStream=net.roseboy.classfinal.JarDecryptor.getInstance().decryptConfigFile(fileName,inputStream,c);");
        aopMap.put("jfinal.line", "62");
    }

    //要加密的jar或war
    private String jarPath = null;
    private B b = new B();
    //要加密的包，多个用逗号隔开
    private List<String> packages = null;
    //-INF/lib下要加密的jar
    private List<String> includeJars = null;
    //排除的类名
    private List<String> excludeClass = null;
    //依赖jar路径
    private List<String> classPath = null;
    //需要加密的配置文件
    private List<String> cfgfiles = null;
    //机器码
    private char[] code = null;

    //jar还是war
    private String jarOrWar = null;
    //工作目录
    private File targetDir = null;
    //agent输出工作目录
    private File targetAgentDir = null;
    //-INF/lib目录
    private File targetLibDir = null;
    //-INF/classes目录
    private File targetClassesDir = null;
    //加密的文件数量
    private Integer encryptFileCount = null;
    //存储解析出来的类名和路径
    private Map<String, String> resolveClassName = new HashMap<>();

    /**
     * 构造方法
     *
     * @param jarPath  要加密的jar或war
     */
    public JarEncryptor(String jarPath) {
        super();
        this.jarPath = jarPath;
    }

    /**
     * 加密jar的主要过程
     *
     * @return 解密后生成的文件的绝对路径
     */
    public String doEncryptJar() {
        if (!jarPath.endsWith(".jar") && !jarPath.endsWith(".war")) {
            throw new RuntimeException("jar/war文件格式有误");
        }
        if (!new File(jarPath).exists()) {
            throw new RuntimeException("文件不存在:" + jarPath);
        }
        Log.debug("机器绑定：" + (StrUtils.isEmpty(this.code) ? "否" : "是"));

        this.jarOrWar = jarPath.substring(jarPath.lastIndexOf(".") + 1);
        Log.debug("加密类型：" + jarOrWar);
        //临时work目录
        this.targetDir = new File(jarPath.replace("." + jarOrWar, Const.LIB_JAR_DIR));
        //临时work目录
        this.targetAgentDir = new File(jarPath.replace("." + jarOrWar, Const.LIB_JAR_DIR+"agent"));

        this.targetLibDir = new File(this.targetDir, ("jar".equals(jarOrWar) ? "BOOT-INF" : "WEB-INF")
                + File.separator + "lib");
        this.targetClassesDir = new File(this.targetDir, ("jar".equals(jarOrWar) ? "BOOT-INF" : "WEB-INF")
                + File.separator + "classes");
        Log.debug("临时目录：" + targetDir);

        //[1]释放所有文件
        long time = System.currentTimeMillis();

//        List<String> allFile = JarUtils.unJar(jarPath, this.targetDir.getAbsolutePath());
        Map<String,byte[]> allFileMap = JarUtils.unZipToMap(jarPath, null);
//        allFile.forEach(s -> Log.debug("释放：" + s));
        //[1.1]内部jar只释放需要加密的jar
        List<String> libJarFiles = new ArrayList<>();
        Map<String,Map<String,byte[]>> unZipJars=new HashMap<>();

        List<String> allFile=new ArrayList<>();
        for (Map.Entry<String,byte[]> stringEntry : allFileMap.entrySet()) {
            String path=stringEntry.getKey();
            byte[] bs=stringEntry.getValue();

            if(bs==null)continue;
            allFile.add(path);
            if (!path.toLowerCase().endsWith(".jar")) {
                continue;
            }
            String name = path.substring(path.lastIndexOf("/") + 1);
            if (StrUtils.isMatchs(this.includeJars, name, false)) {
//                String targetPath = path.substring(0, path.length() - 4) + Const.LIB_JAR_DIR;
                Map<String,byte[]> inMap=JarUtils.unZipToMap(bs,null);
                unZipJars.put(path,inMap);
//                List<String> files = JarUtils.unJar(path, targetPath);
//                files.forEach(s -> Log.debug("释放：" + s));
                libJarFiles.add(path);
//                libJarFiles.addAll(files);
            }

        }

        for (Map.Entry<String, Map<String, byte[]>> content : unZipJars.entrySet()) {
           List<String> encodeFiles=encodeMap(content.getValue());
            System.out.println("加密内部jar包"+content.getKey()+":"+(System.currentTimeMillis()-time));
            time = System.currentTimeMillis();
//            将jar包内加密类迁移到外部包根目录，方便程序读取解析
            for (String encodeFile : encodeFiles) {
                String path="META-INF/"+Const.FILE_NAME+"/"+encodeFile;
                allFileMap.put(path,content.getValue().get(path));
                content.getValue().remove(path);
            }
        }


        allFile.addAll(libJarFiles);
        System.out.println("释放所有文件:"+(System.currentTimeMillis()-time));
        time = System.currentTimeMillis();

        //压缩静态文件
//        allFile.forEach(s -> {
//            if (!s.endsWith(".ftl")) {
//                return;
//            }
//            File file = new File(s);
//            String code = IoUtils.readTxtFile(file);
//            code = HtmlUtils.removeComments(code);
//            code = HtmlUtils.removeBlankLine(code);
//            IoUtils.writeTxtFile(file, code);
//        });

        //[2]提取所有需要加密的class文件
//        List<File> classFiles = filterClasses(allFile);
        List<String> classFiles = getClasses(allFile);
        System.out.println("提取所有需要加密的class文件:"+classFiles.size()+":"+(System.currentTimeMillis()-time));
        time = System.currentTimeMillis();

        //[3]将本项目的代码添加至jar中
//        addClassFinalAgent();

        //[4]将正常的class加密，压缩另存
        List<String> encryptClass = encryptClass(classFiles,allFileMap);
        this.encryptFileCount = encryptClass.size();
        System.out.println("将正常的class加密，压缩另存:"+(System.currentTimeMillis()-time));
        time = System.currentTimeMillis();

        //[5]清空class方法体，并保存文件
        clearClassMethod(classFiles,allFileMap);
        System.out.println("清空class方法体，并保存文件:"+(System.currentTimeMillis()-time));
        time = System.currentTimeMillis();

        //[6]加密配置文件
        encryptConfigFile();
        System.out.println("加密配置文件:"+(System.currentTimeMillis()-time));
        time = System.currentTimeMillis();

        //[7]打包回去
        String result = packageJar(allFileMap,unZipJars);

        System.out.println("打包回去:"+(System.currentTimeMillis()-time));
        return result;
    }

    private List<String> encodeMap(Map<String, byte[]> content) {
        List<String> files=new ArrayList<>();
        for (Map.Entry<String, byte[]> entry : content.entrySet()) {
            if(entry.getValue()!=null)files.add(entry.getKey());
        }
        List<String> classFiles = getClasses(files);
       List<String> result= encryptClass(classFiles,content);
        clearClassMethod(classFiles,content);
        return result;
    }

    /**
     * 找出所有需要加密的class文件
     *
     * @param allFile 所有文件
     * @return 待加密的class列表
     */
    public List<String> getClasses(List<String> allFile) {
        List<String> classFiles = new ArrayList<>();
        for (String file : allFile) {
            if (!file.endsWith(".class")) {
                continue;
            }
            //解析出类全名
            String className = resolveClassName(file, true);
            //判断包名相同和是否排除的类
            if (StrUtils.isMatchs(this.packages, className, false)
                    && !StrUtils.isMatchs(this.excludeClass, className, false)) {
                classFiles.add(file);
                Log.debug("待加密: " + file);
            }

        }
        return classFiles;
    }

    /**
     * 加密class文件，放在META-INF/classes里
     *
     * @param classFiles jar/war 下需要加密的class文件
     * @return 已经加密的类名
     */
    private List<String> encryptClass(List<File> classFiles) {
        List<String> encryptClasses = new ArrayList<>();

        //加密后存储的位置
        File metaDir = new File(this.targetDir, "META-INF" + File.separator + Const.FILE_NAME);
        if (!metaDir.exists()) {
            metaDir.mkdirs();
        }



        //有机器码
        if (StrUtils.isNotEmpty(this.code)) {
            File configCode = new File(metaDir, Const.CONFIG_CODE);
            IoUtils.writeFile(configCode, toBytes(EncryptUtils.md5(this.code)));
        }

        //加密另存
        for (File classFile : classFiles) {

            String className = classFile.getName();
            if (className.endsWith(".class")) {
                className = resolveClassName(classFile.getAbsolutePath(), true);
            }
            byte[] bytes = IoUtils.readFileToByte(classFile);
            bytes = b.e(bytes, className.getBytes(StandardCharsets.UTF_8));
            //有机器码，再用机器码加密一遍
            if (StrUtils.isNotEmpty(this.code)) {
                bytes = b.e(bytes, toBytes(this.code));
            }
            File targetFile = new File(metaDir, className);
            IoUtils.writeFile(targetFile, bytes);
            encryptClasses.add(className);
            Log.debug("加密：" + className);

        }

        return encryptClasses;
    }
    /**
     * 加密class文件，放在META-INF/classes里
     *
     * @param classFiles jar/war 下需要加密的class文件
     * @return 已经加密的类名
     */
    private List<String> encryptClass(List<String> classFiles,Map<String,byte[]> map) {
        List<String> encryptClasses = new ArrayList<>();

        String newPathRoot="META-INF/"+Const.FILE_NAME+"/";
        //加密后存储的位置
//        File metaDir = new File(this.targetDir, "META-INF" + File.separator + Const.FILE_NAME);
//        if (!metaDir.exists()) {
//            metaDir.mkdirs();
//        }


        //有机器码
        if (StrUtils.isNotEmpty(this.code)) {
            map.put(newPathRoot+Const.CONFIG_CODE, toBytes(EncryptUtils.md5(this.code)));
        }

        //加密另存
        for (String className : classFiles) {

            byte[] bytes = map.get(className);
            if (className.endsWith(".class")) {
                className = resolveClassName(className, true);
            }
            bytes = b.e(bytes, className.getBytes(StandardCharsets.UTF_8));
//            byte[] bs = b.e(bytes, className.getBytes(StandardCharsets.UTF_8));
//            C.init();
//            byte[] bds = b.d(bs, className.getBytes(StandardCharsets.UTF_8));
//            if(Arrays.equals(bds,bytes)){
//                System.out.println("加密成功");
//                bytes=bs;
//            }else{
//                System.out.println("加密失败");
//            }

            //有机器码，再用机器码加密一遍
            if (StrUtils.isNotEmpty(this.code)) {
                bytes = b.e(bytes, toBytes(this.code));
            }
//            File targetFile = new File(metaDir, className);
//            IoUtils.writeFile(targetFile, bytes);
            map.put(newPathRoot+className,bytes);

            encryptClasses.add(className);
            Log.debug("加密：" + className);

        }

        map.put(newPathRoot+Const.CONFIG_PASSHASH, b.e(new byte[]{}, new byte[]{}));
        return encryptClasses;
    }


    /**
     * 清空class文件的方法体，并保留参数信息
     *
     * @param classFiles jar/war 下需要加密的class文件
     */
    private void clearClassMethod(List<String> classFiles,Map<String,byte[]> map) {
        //初始化javassist
        ClassPool pool = ClassPool.getDefault();
        Map<String, CtClass> ctClassMap=new HashMap<>();

        for (Map.Entry<String, byte[]> entry : map.entrySet()) {
            String classFile=entry.getKey();
            byte[] bs=entry.getValue();

            if(bs==null)continue;
            if(classFile.endsWith("class")){
                try {
                    ctClassMap.put(classFile, pool.makeClass(new ByteArrayInputStream(bs)));
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        AtomicInteger successCount= new AtomicInteger();
        //[2]修改class方法体，并保存文件
        ClassUtils.errSet.clear();
        for (String classFile : classFiles) {

            //解析出类全名
            String className = resolveClassName(classFile, true);
            byte[] bts = null;
            try {
                Log.debug("清除方法体: " + className);
                bts = ClassUtils.rewriteAllMethods(pool, ctClassMap.get(classFile));
                successCount.getAndIncrement();
            } catch (Exception e) {
                e.printStackTrace();
                Log.debug("ERROR:" + e.getMessage());
            }
            if (bts != null) {
                map.put(classFile,bts);
            }

        }
        System.out.println("ClassUtils.errSet:"+ClassUtils.errSet);
        System.out.println("清除数量："+successCount.get()+"，总数："+classFiles.size());
    }


    /**
     * 加密classes下的配置文件
     */
    private void encryptConfigFile() {
        if (this.cfgfiles == null || this.cfgfiles.size() == 0) {
            return;
        }

        //支持的框架
        //String[] supportFrame = {"spring", "jfinal"};
        String[] supportFrame = {"spring"};
        //需要注入解密功能的class
        List<File> aopClass = new ArrayList<>(supportFrame.length);

        // [1].读取配置文件时解密
        for (String name : supportFrame) {
            String javaCode = aopMap.get(name + ".code");
            String clazz = aopMap.get(name + ".class");
            Integer line = Integer.parseInt(aopMap.get(name + ".line"));
            javaCode = javaCode.replace("${passchar}", StrUtils.toCharArrayCode(new char[]{}));
            byte[] bytes = null;
            try {
                String thisJar = this.getClass().getProtectionDomain().getCodeSource().getLocation().getPath();
                //获取 框架 读取 配置文件的类,将密码注入该类
                bytes = ClassUtils.insertCode(clazz, javaCode, line, this.targetLibDir, new File(thisJar));
            } catch (Exception e) {
                e.printStackTrace();
                Log.debug(e.getClass().getName() + ":" + e.getMessage());
            }
            if (bytes != null) {
                File cls = new File(this.targetDir, clazz.split("#")[0] + ".class");
                IoUtils.writeFile(cls, bytes);
                aopClass.add(cls);
            }

        }

        //加密读取配置文件的类
        this.encryptClass(aopClass);
        for (File cls : aopClass) {
            cls.delete();
        }


        //[2].加密配置文件
        List<File> configFiles = new ArrayList<>();
        File[] files = this.targetClassesDir.listFiles();
        if (files == null) {
            return;
        }
        for (File file : files) {
            if (file.isFile() && StrUtils.isMatchs(this.cfgfiles, file.getName(), false)) {
                configFiles.add(file);
            }
        }
        //加密
        this.encryptClass(configFiles);
        //清空
        for (File file : configFiles) {
            IoUtils.writeTxtFile(file, "");
        }
    }

    /**
     * 压缩成jar
     *
     * @return 打包后的jar绝对路径
     */
    private String packageJar(Map<String,byte[]> map,Map<String,Map<String,byte[]>> jarMap) {
        //[1]先打包lib下的jar

        for (Map.Entry<String, Map<String,byte[]>> stringMapEntry : jarMap.entrySet()) {
            map.put(stringMapEntry.getKey(),JarUtils.doJar(stringMapEntry.getValue()));
        }
        //[2]再打包jar
        String targetJar = jarPath.replace("." + jarOrWar, "-encrypted." + jarOrWar);
        String result = JarUtils.doJar(map, targetJar);
        Log.debug("打包: " + targetJar);
        return result;
    }

    /**
     * 根据class的绝对路径解析出class名称或class包所在的路径
     *
     * @param fileName    class绝对路径
     * @param classOrPath true|false
     * @return class名称|包所在的路径
     */
    private String resolveClassName(String fileName, boolean classOrPath) {
        String result = resolveClassName.get(fileName + classOrPath);
        if (result != null) {
            return result;
        }
        String file = fileName.substring(0, fileName.length() - 6);
        String K_CLASSES = File.separator + "classes" + File.separator;
        String K_CLASSES2 = "/classes/";
        String K_LIB = File.separator + "lib" + File.separator;
        String K_LIB2 = "/lib/";

        String clsPath;
        String clsName;
        //lib内的的jar包
        if (file.contains(K_LIB)) {
            clsName = file.substring(file.indexOf(Const.LIB_JAR_DIR, file.indexOf(K_LIB))
                    + Const.LIB_JAR_DIR.length() + 1);
            clsPath = file.substring(0, file.length() - clsName.length() - 1);
        }
        else if (file.contains(K_LIB2)) {
            clsName = file.substring(file.indexOf(Const.LIB_JAR_DIR, file.indexOf(K_LIB2))
                    + Const.LIB_JAR_DIR.length() + 1);
            clsPath = file.substring(0, file.length() - clsName.length() - 1);

        }
        //jar/war包-INF/classes下的class文件
        else if (file.contains(K_CLASSES)) {
            clsName = file.substring(file.indexOf(K_CLASSES) + K_CLASSES.length());
            clsPath = file.substring(0, file.length() - clsName.length() - 1);

        }else if (file.contains(K_CLASSES2)) {
            clsName = file.substring(file.indexOf(K_CLASSES2) + K_CLASSES2.length());
            clsPath = file.substring(0, file.length() - clsName.length() - 1);

        }
        //jar包下的class文件
        else {
            clsName = file;
            clsPath = file;
        }
        result = classOrPath ? clsName.replace(File.separator, ".").replace("/", ".") : clsPath;
        resolveClassName.put(fileName + classOrPath, result);
        return result;
    }


    public static void main(String[] args) {
        String file ="C:\\Program Files\\Java\\jdk1.8.0_291\\bin\\java.exe";


        System.out.println(new File(".").getParent());
    }
    public Integer getEncryptFileCount() {
        return encryptFileCount;
    }

    public void setPackages(List<String> packages) {
        this.packages = packages;
    }

    public void setIncludeJars(List<String> includeJars) {
        this.includeJars = includeJars;
    }

    public void setExcludeClass(List<String> excludeClass) {
        this.excludeClass = excludeClass;
    }

    public void setClassPath(List<String> classPath) {
        this.classPath = classPath;
    }

    public void setCfgfiles(List<String> cfgfiles) {
        this.cfgfiles = cfgfiles;
    }

    public void setCode(char[] code) {
        this.code = code;
    }

}
